{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "nbsphinx": "hidden"
   },
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Seasonal trends with the GOZCARDS Dataset\n",
    "\n",
    "Here we calculate seasonal trends using the GOZCARDS dataset by regressing to the VMR monthly zonal means using seasonal terms in our predictors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import xarray as xr\n",
    "import numpy as np\n",
    "from LOTUS_regression.regression import regress_all_bins\n",
    "from LOTUS_regression.predictors.seasonal import add_seasonal_components\n",
    "from LOTUS_regression.predictors import load_data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The GOZCARDS data is in multiple NetCDF4 files.  Load them all in and concatenate on the time dimension.  Also only take values in the latitude range -60 to 60."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "GOZCARDS_FILES = r'/media/users/data/LOTUS/GOZCARDS/*.nc4'\n",
    "\n",
    "data = xr.decode_cf(xr.open_mfdataset(GOZCARDS_FILES, concat_dim='time', group='Merged', engine='netcdf4'))\n",
    "\n",
    "data = data.loc[dict(lat=slice(-60, 60))]\n",
    "\n",
    "print(data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load some standard predictors and add a constant"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "predictors = load_data('pred_baseline_pwlt.csv')\n",
    "\n",
    "print(predictors.columns)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Currently our predictors have no seasonal dependence.  Add in some seasonal terms with different numbers of Fourier components."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "predictors = add_seasonal_components(predictors, {'constant': 4, 'linear_pre': 4, 'linear_post': 4, 'qboA': 2, 'qboB': 2})\n",
    "\n",
    "print(predictors[:5])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Perform the regression and convert the coefficients to percent anomaly.  We pass include_monthly_fits = True so that\n",
    "the seasonal fits are used to calculate monthly trends.  The results at the end will include 'linear_post_monthly'\n",
    "and 'linear_post_monthly_std' that are the monthly trend terms and errors respectively"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "results = regress_all_bins(predictors, data['average'], tolerance=0.1, include_monthly_fits=True)\n",
    "\n",
    "# Convert to ~ percent\n",
    "results /= data['average'].mean(dim='time')\n",
    "results *= 100\n",
    "\n",
    "print(results)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import LOTUS_regression.plotting.trends as trends\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# Plot the seasonal post trends at the level closest to 2 hPa (2.15 hPa)\n",
    "trends.plot_with_confidence(results.sel(lev=2, method='nearest'), 'linear_post_monthly', x='lat', y='month')\n",
    "plt.xlabel('Latitude')\n",
    "plt.ylabel('Month')"
   ]
  }
 ],
 "metadata": {
  "celltoolbar": "Edit Metadata",
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}